#include "Lattice.h"

const int cx[Q_DIM] = 	{0, 1, 0, -1,  0, 1, -1, -1,  1};
const int cy[Q_DIM] = 	{0, 0, 1,  0, -1, 1,  1, -1, -1};
const int opDir[Q_DIM]= {0, 3, 4,  1,  2, 7,  8,  5,  6};
const double w[Q_DIM] = {4./9., 1./9., 1./9., 1./9., 1./9., 1./36., 1./36., 1./36., 1./36.};

Lattice::Lattice() {
}

// void Lattice::relaxRho() {
// 	array2< vec<double> > u0;
// 	u0.Allocate(dimX, dimY, sizeof(vec<double>));
// 
// 	//inicializacion de rho y u
// 	for (int i = 0; i < dimX; i++)
// 		for (int j = 0; j < dimY; j++) {
// 			double u0x = -0.05*cos(i*2.0*M_PI/dimX)*sin(j*2.0*M_PI/dimY);
// 			double u0y = 0.05*cos(j*2.0*M_PI/dimY)*sin(i*2.0*M_PI/dimX);
// 			u0[i][j] = {u0x, u0y};
// 			grid[i][j].ux = u0[i][j].x;
// 			grid[i][j].uy = u0[i][j].y;
// 
// 			grid[i][j].r = 1.0;
// 			grid[i][j].dr = 0.0;
// 
// 			grid[i][j].m[iDR] = 0.0;
// 			grid[i][j].m[iJX] = u0x*grid[i][j].r; //esto está bien?
// 			grid[i][j].m[iJY] = u0y*grid[i][j].r;
// 			grid[i][j].m[iE] = -2.0*alpha2*grid[i][j].m[iDR] + 3.0*(pow(grid[i][j].m[iJX],2) + pow(grid[i][j].m[iJY],2));
// 			grid[i][j].m[iEPSILON] = grid[i][j].m[iDR] - 3.0*(pow(grid[i][j].m[iJX],2) + pow(grid[i][j].m[iJY],2));
// 			grid[i][j].m[iQX] = -grid[i][j].m[iJX];
// 			grid[i][j].m[iQY] = -grid[i][j].m[iJY];
// 			grid[i][j].m[iPXX] = pow(grid[i][j].m[iJX],2) - pow(grid[i][j].m[iJY],2);
// 			grid[i][j].m[iPXY] = grid[i][j].m[iJX]*grid[i][j].m[iJY];
// 		}
// 
// 	double rhoDiff = 1000;
// 	while (rhoDiff > 1e-08) {
// 		rhoDiff = 0.0;
// 
// 		for (int i = 0; i < dimX; i++)
// 			for (int j = 0; j < dimY; j++) {
// 				grid[i][j].m[iJX] = grid[i][j].m[iJX] - 1.0*(grid[i][j].m[iJX] - u0[i][j].x);
// 				grid[i][j].m[iJY] = grid[i][j].m[iJY] - 1.0*(grid[i][j].m[iJY] - u0[i][j].y);
// 
// 				double _Eeq = -2.0*alpha2*grid[i][j].m[iDR] + 3.0*(pow(grid[i][j].m[iJX],2) + pow(grid[i][j].m[iJY],2));
// 				double _EPSILONeq = grid[i][j].m[iDR] - 3.0*(pow(grid[i][j].m[iJX],2) + pow(grid[i][j].m[iJY],2));
// 				double _QXeq = -grid[i][j].m[iJX];
// 				double _QYeq = -grid[i][j].m[iJY];
// 				double _PXXeq = pow(grid[i][j].m[iJX],2) - pow(grid[i][j].m[iJY],2);
// 				double _PXYeq = grid[i][j].m[iJX]*grid[i][j].m[iJY];
// 
// 				grid[i][j].m[iE] = 		grid[i][j].m[iE] - 		grid[i][j].s[iE]*(grid[i][j].m[iE] - 				_Eeq);
// 				grid[i][j].m[iEPSILON] = grid[i][j].m[iEPSILON] - grid[i][j].s[iEPSILON]*(grid[i][j].m[iEPSILON] - 	_EPSILONeq);
// 				grid[i][j].m[iQX] = 		grid[i][j].m[iQX] - 		grid[i][j].s[iQX]*(grid[i][j].m[iQX] - 			_QXeq);
// 				grid[i][j].m[iQY] = 		grid[i][j].m[iQY] - 		grid[i][j].s[iQY]*(grid[i][j].m[iQY] - 			_QYeq);
// 				grid[i][j].m[iPXX] = 	grid[i][j].m[iPXX] - 	grid[i][j].s[iPXX]*(grid[i][j].m[iPXX] - 			_PXXeq);
// 				grid[i][j].m[iPXY] = 	grid[i][j].m[iPXY] - 	grid[i][j].s[iPXY]*(grid[i][j].m[iPXY] - 			_PXYeq);
// 				grid[i][j].calcDF();
// 			}
// 
// 		for (int i = 0; i < dimX; i++)
// 			for (int j = 0; j < dimY; j++) {
// 				double _dr = grid[i][j].dr;
// 				grid[i][j].updateRhoAndU();
// 				rhoDiff += abs(_dr - grid[i][j].dr);
// 			}
// 
// 		for (int i = 0; i < dimX; i++)
// 			for (int j = 0; j < dimY; j++)
// 				streamCell(i, j);
// 		swapGrids();
// 
// 		for (int i = 0; i < dimX; i++)
// 			for (int j = 0; j < dimY; j++)
// 				grid[i][j].calcMomenta();
// 
// 		//cout << "pop = " << getPopulation() << endl;
// 		cout << rhoDiff << endl;
// 	}
// 		for (int i = 0; i < dimX; i++)
// 			for (int j = 0; j < dimY; j++) {
// 				grid[i][j].calcMomenta();
// 				grid[i][j].dr = grid[i][j].m[iDR];
// 				grid[i][j].calcDF();
// 			}
// }

void Lattice::relaxRho() {
	vector<vector<double> > u0x(dimY, vector<double>(dimX, 0));
	vector<vector<double> > u0y(dimY, vector<double>(dimX, 0));
	vector<vector<double> > drOld(dimY, vector<double>(dimX, 0));

	for (int i = 0; i < dimX; i++)
		for (int j = 0; j < dimY; j++) {
			u0x[i][j] = grid[i][j].ux;
			u0y[i][j] = grid[i][j].uy;
			grid[i][j].r = 1.0;
			grid[i][j].dr = 0.0;
			grid[i][j].m[iDR] = 0.0;
		}

	double rhoDiff = 1000;
	while (rhoDiff > 3e-16) {
	//for(int rep=1; rep<46; rep++) {
		for (int i = 0; i < dimX; i++)
			for (int j = 0; j < dimY; j++)
				grid[i][j].collide(u0x[i][j], u0y[i][j]); //calculo momentos, fijo u, relajo, calculo DF, updateRhoAndU

		for (int i = 0; i < dimX; i++)
			for (int j = 0; j < dimY; j++) {
				//grid[i][j].updateRhoAndU();
				drOld[i][j] = grid[i][j].dr;
			}

		for (int i = 0; i < dimX; i++) //obtengo f nuevos en gridNew
			for (int j = 0; j < dimY; j++)
				streamCell(i, j);

		for (int i = 0; i < dimX; i++)
			for (int j = 0; j < dimY; j++)
				for (int l = 0; l < Q_DIM; l++)
					grid[i][j].f[l] = gridNew[i][j].f[l]; //paso los f nuevos a grid
				
		rhoDiff = 0.0;
		for (int i = 0; i < dimX; i++)
			for (int j = 0; j < dimY; j++) {
				grid[i][j].updateRhoAndU();
				rhoDiff += abs(drOld[i][j] - grid[i][j].dr);
			}
		cout << rhoDiff << endl;
	}
}

void Lattice::init(int nX, int nY, double _rho, double _se, double _seps, double _sq, double _snu) {
	dimX = nX;
	dimY = nY;

	grid = vector< vector<Cell> >(dimX, vector<Cell>(dimY));
	gridNew = vector< vector<Cell> >(dimX, vector<Cell>(dimY));

	for (int i = 0; i < dimX; i++)
		for (int j = 0; j < dimY; j++) {
			grid[i][j].s[iDR] = 0.0;
			grid[i][j].s[iE] = _se;
			grid[i][j].s[iEPSILON] = _seps;
			grid[i][j].s[iJX] = 0.0;
			grid[i][j].s[iQX] = _sq;
			grid[i][j].s[iJY] = 0.0;
			grid[i][j].s[iQY] = _sq;
			grid[i][j].s[iPXX] = _snu;
			grid[i][j].s[iPXY] = _snu;
		}
	resetGrid(_rho);
}

void Lattice::resetGrid(double _rho) {
	for (int i = 0; i < dimX; i++)
		for (int j = 0; j < dimY; j++)
			grid[i][j].setU(0.0, 0.0, _rho);
	//*gridAux = *grid;
}

void Lattice::streamCell(int i, int j) {
	for (int l = 0; l < Q_DIM; l++) {
		int fromCellX = i + cx[opDir[l]];
		int fromCellY = j + cy[opDir[l]];

		if (fromCellX < 0)
			fromCellX = dimX-1;
		else if (fromCellX == dimX)
			fromCellX = 0;

		if (fromCellY < 0)
			fromCellY = dimY-1;
		else if (fromCellY == dimY)
			fromCellY = 0;

		gridNew[i][j].f[l] = grid[fromCellX][fromCellY].f[l];
	}
}

void Lattice::update() {
	for (int i = 0; i < dimX; i++)
		for (int j = 0; j < dimY; j++)
			streamCell(i, j);

	for (int i = 0; i < dimX; i++)
		for (int j = 0; j < dimY; j++)
			for (int l = 0; l < Q_DIM; l++)
				grid[i][j].f[l] = gridNew[i][j].f[l];
		
	for (int i = 0; i < dimX; i++)
		for (int j = 0; j < dimY; j++)
			grid[i][j].collide();
}

double Lattice::energy() {
	double ret = 0.0;
	for(int i=0; i < dimX; i++)
		for(int j=0; j < dimY; j++) {
			double _ux = grid[i][j].ux;
			double _uy = grid[i][j].uy;
			double _rho = grid[i][j].r;// + grid[i][j].dr;
			//ret += _rho*(_ux*_ux + _uy*_uy);
			ret += (_ux*_ux + _uy*_uy);
		}
	return ret/2.0;
}

double Lattice::population() {
	double ret = 0.0;
	for (int i = 0; i < dimX; i++)
		for (int j = 0; j < dimY; j++)
			ret += grid[i][j].r + grid[i][j].dr;
	return ret;
}

double Lattice::divergence() {
	double _rmsdiv = 0.0;
	for (int i = 1; i < dimX-1; i++)
		for (int j = 1; j < dimY-1; j++) {
			double _dx = (grid[i+1][j].ux - grid[i-1][j].ux)/2.0;
			double _dy = (grid[i][j+1].uy - grid[i][j-1].uy)/2.0;
			double _div = _dx + _dy;
			_rmsdiv += abs(_div);
		}
	return _rmsdiv/((dimX-2)*(dimY-2));
}

void Lattice::setU(int i, int j, double _ux, double _uy) {
	if ((i >= 0) && (i < dimX) && (j >= 0) && (j < dimY))
		grid[i][j].setU(_ux, _uy, grid[i][j].r);
}

void Lattice::setU(int i, int j, double _ux, double _uy, double _rho) {
	if ((i >= 0) && (i < dimX) && (j >= 0) && (j < dimY))
		grid[i][j].setU(_ux, _uy, _rho);
}
